# uamcf/methods/__init__.py
from .base import CompressionMethod
from .quantization import Quantization
from .pruning import Pruning
from .distillation import Distillation
from .factorization import Factorization
from .low_rank import LowRankDecomposition
from .weight_sharing import WeightSharing
from .parameter_sharing import ParameterSharing
from .nas import NeuralArchitectureSearch

# 注册所有压缩方法
METHODS_REGISTRY = {
    "quantization": Quantization,
    "pruning": Pruning,
    "distillation": Distillation,
    "factorization": Factorization,
    "low_rank": LowRankDecomposition,
    "weight_sharing": WeightSharing,
    "parameter_sharing": ParameterSharing,
    "nas": NeuralArchitectureSearch
}

def get_method(method_name):
    """获取压缩方法类"""
    return METHODS_REGISTRY.get(method_name)

def get_method_registry():
    """获取所有已注册的压缩方法"""
    return METHODS_REGISTRY

def register_method(name, method_cls):
    """注册新的压缩方法"""
    METHODS_REGISTRY[name] = method_cls

__all__ = [
    'CompressionMethod',
    'Quantization',
    'Pruning',
    'Distillation',
    'Factorization',
    'LowRankDecomposition',
    'WeightSharing',
    'ParameterSharing',
    'NeuralArchitectureSearch',
    'get_method',
    'get_method_registry',
    'register_method'
]